namespace IBS.Math

module Forecasting = 
    open System

    let internal indToValues (card: _[]) ind =
        let max = Seq.reduce ( * ) card
        let len = card.Length
        if ind >= max || ind < 0 then failwith "index out of range"
        Array.init len (fun i -> 
            let rep = Seq.take (i+1) card |> Seq.reduce ( * )
            if i = 0 then ind % card.[i] else ind % rep / (rep / card.[i]))

    let internal getValues (card: _[]) =
        let size = Seq.reduce ( * ) card
        Array.init size (indToValues card) 

    let internal teq (x, y) = x = y
    let tmap f (x, y) = f x, f y

    let internal mapToInds (card: _[]) =
        let all = getValues card
        fun x -> Seq.findIndex (fun a -> Seq.zip x a |> Seq.forall teq) all

    let internal findVals xVars yVars (xs: _[]) =
        let xvs = Set.map fst xVars
        let m = Seq.map (fun (yv, _) -> xvs.Contains yv) yVars |> Seq.toArray
        let allYs = getValues (Seq.map snd yVars |> Seq.toArray)
        let xasmts = mapToInds (Seq.map snd xVars |> Seq.toArray)
        Seq.map (fun ys -> 
            let res = Seq.zip ys m |> Seq.filter snd |> Seq.map fst |> Seq.toArray
            xs.[xasmts res]) allYs

    let internal sumAllYBy asmt xVars yVars (yVals: _[]) =
        let xvs = Set.map fst xVars
        let m = Seq.map (fun (yv, _) -> xvs.Contains yv) yVars |> Seq.toArray
        let allYs = getValues (Seq.map snd yVars |> Seq.toArray)
        Seq.mapi (fun i ys -> 
            let res = Seq.zip ys m |> Seq.filter snd |> Seq.map fst |> Seq.zip asmt |> Seq.forall (fun (x, y) -> x = y)
            if res then yVals.[i] else 0.) allYs   
        |> Seq.reduce (+)
    
    type Factor = {
        Vars: (int * int) Set // # of var and its cardinality
        Val : float []
    }
    with 
        static member ( * ) (x, y) = 
            let varsZ = x.Vars + y.Vars
            if Seq.countBy fst varsZ |> Seq.exists (snd >> ((<) 1)) then
                failwith "cardinalities don't match"
            let xas, yas = tmap (fun v -> mapToInds (Seq.map snd v.Vars |> Seq.toArray)) (x, y)
            let xVals = findVals x.Vars varsZ x.Val
            let yVals = findVals y.Vars varsZ y.Val
            
            { Vars = varsZ; Val = Seq.zip xVals yVals |> Seq.map (fun (v1, v2) -> v1 * v2) |> Seq.toArray }
         
        member x.Marginalize vs =
            let vars = Set.filter (fun (x, _) -> Seq.exists ((=)x) vs |> not) x.Vars
            let card = Seq.map snd vars |> Seq.toArray
            let vals = Array.map (fun ai ->
                sumAllYBy ai vars x.Vars x.Val) (getValues card)
            { Vars = vars; Val = vals }
            
        member x.ApplyEvidence (vs: _ Set) =
            let vs = Seq.filter (fun (var, _) -> Seq.exists (fun (xvar, _) -> var = xvar) x.Vars) vs |> Seq.toArray
            let vars = Array.map fst vs
            let vals = Array.map snd vs
            let m = Seq.map (fun (xv, _) -> Array.IndexOf(vars, xv) <> -1) x.Vars |> Seq.toArray
            let res = Array.mapi (fun i asmt ->
                let res = Seq.zip asmt m |> Seq.filter snd |> Seq.map fst |> Seq.zip vals |> Seq.forall (fun (x, y) -> x = y)
                if res then x.Val.[i] else 0.) (getValues (Seq.map snd x.Vars |> Seq.toArray))
            { x with Val = res }
            
        static member ComputeMarginal vs (fs: #seq<Factor>) e = 
            let fj = Seq.reduce ( * ) fs
            let fe = fj.ApplyEvidence e
            let sum = Seq.sum fe.Val
            let feval = Array.map (fun v -> v / sum) fe.Val //normalization
            let vars = Set.map fst fe.Vars - set vs
            { fe with Val = feval }.Marginalize vars
 
    open System.ComponentModel

    type Age =         
        | [<Description "Below 22">] Below22 = 0        
        | [<Description "Between 22 and 50">] Between22And50 = 1
        | [<Description "Above 50">] Above50 = 2
    
    type DebtRatio =  // Debt / Income
        | High = 0
        | Low = 1
            
    // depends on Age & DebtRatio
    type PaymentHistory = 
        | Excellent  = 0
        | Acceptable = 1
        | Unacceptable = 2
      
    // depends on Age and PaymentHistory  
    type Reliability =
        | Reliable = 0
        | Unreliable = 1
    
    // depends on DebtRatio
    type Income =
        | High = 0
        | Medium = 1
        | Low = 2
        
    // depends on Income
    type Assets =
        | High = 0
        | Medium = 1
        | Low = 2
        
    // depends on Assets and Income
    type FutureIncome =
        | Promising = 0
        | [<Description "Not Promising">] NotPromising = 1
    
    // depends on DebtRatio, Reliability and FutureIncome
    type Forecast =
        | Positive = 0
        | Negative = 1 

    let inline iv (p: 'a option) = Option.bind (fun x -> Some (int x)) p
    let inline toOption (p: Nullable<'a>) = if p.HasValue then Some (int p.Value) else None

    // model for predictions and c# compatibility
    type Model() =
        let fAge = { Vars = set [1, 3]; Val = [| 0.33;0.4;0.27 |] }
        let fRatio = { Vars = set [2, 2]; Val = [| 0.5;0.5 |] }
        let fHistory = { Vars = set [1, 3; 2, 2; 3, 3]; Val = [| 0.3;0.5;0.6;0.2;0.3;0.4;0.4;0.35;0.3;0.3;0.25;0.2;0.3;0.15;0.1;0.5;0.45;0.4 |] }
        let fReliability = { Vars = set [1, 3; 3, 3; 4, 2]; Val = [| 0.6;0.8;0.95;0.5;0.6;0.7;0.15;0.2;0.3;0.4;0.2;0.05;0.5;0.4;0.3;0.85;0.8;0.7 |] }
        let fIncome = { Vars = set [2, 2; 5, 3]; Val = [| 0.3;0.7;0.4;0.2;0.3;0.1 |] }
        let fAssets = { Vars = set [5, 3; 6, 3]; Val = [| 0.85;0.6;0.3;0.1;0.3;0.4;0.05;0.1;0.3 |] }
        let fFutureIncome = { Vars = set [5, 3; 6, 3; 7, 2]; Val = [| 0.9;0.8;0.7;0.7;0.6;0.5;0.5;0.3;0.2;0.1;0.2;0.3;0.3;0.4;0.5;0.5;0.7;0.8 |] }
        let fForecast = { Vars = set [2, 2; 4, 2; 7, 2; 8, 2]; Val = [| 0.8;0.9;0.4;0.5;0.5;0.7;0.1;0.2;0.2;0.1;0.6;0.5;0.5;0.3;0.9;0.8 |] }
        let factors = [fAge; fRatio; fHistory; fReliability; fIncome; fAssets; fFutureIncome; fForecast]   
        
        member x.MarginalizeOver vars ps =
            Factor.ComputeMarginal vars factors (set ps)
      
        member x.Forecast(?age: Age, ?ratio: DebtRatio, ?history: PaymentHistory, ?reliability: Reliability, 
                          ?income: Income, ?assets: Assets, ?futureIncome: FutureIncome) =
            let ps = [|iv age; iv ratio; iv history; iv reliability; iv income; iv assets; iv futureIncome|]
                     |> Array.mapi (fun i x -> (i + 1), x)
                     |> Array.filter (snd >> Option.isSome)
                     |> Array.map (fun (i, Some x) -> i, x)

            let marginal = x.MarginalizeOver [8] ps
            Seq.zip marginal.Val [Forecast.Positive; Forecast.Negative] |> Seq.toArray

        static member Forecast(age: Nullable<Age>, ratio: Nullable<DebtRatio>, history: Nullable<PaymentHistory>, reliability: Nullable<Reliability>,
                               income: Nullable<Income>, assets: Nullable<Assets>, futureIncome: Nullable<FutureIncome>): float =
            let ps = [|toOption age; toOption ratio; toOption history; toOption reliability; toOption income; toOption assets; toOption futureIncome|]
                     |> Array.mapi (fun i x -> (i + 1), x)
                     |> Array.filter (snd >> Option.isSome)
                     |> Array.map (fun (i, Some x) -> i, x)
            let marginal = Model().MarginalizeOver [8] ps
            marginal.Val.[0]
            

